"""
conversation_manager.py
对话管理功能的实现
"""

import json
import os
from datetime import datetime
import re
from PyQt6.QtWidgets import QMessageBox, QInputDialog, QFileDialog, QLineEdit
from PyQt6.QtCore import QObject, pyqtSignal

class ConversationManager(QObject):
    """对话管理器类"""
    
    # 定义信号
    conversation_updated = pyqtSignal()  # 当对话列表更新时发出
    error_occurred = pyqtSignal(str)     # 当发生错误时发出
    
    MAX_CONVERSATIONS = 30  # 最大对话数量
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.chat_history = {}           # 格式: {model_name: {conversation_id: [messages]}}
        self.current_conversation = {}    # 格式: {model_name: conversation_id}
        self.conversation_names = {}      # 格式: {model_name: {conversation_id: name}}
        self.next_conversation_id = {}    # 格式: {model_name: next_id}
        self.load_chat_history()         # 加载历史对话
        
    def create_new_conversation(self, model_name):
        """创建新对话"""
        if not model_name:
            self.error_occurred.emit("请先选择一个模型！")
            return None
            
        # 检查对话数量
        if model_name in self.chat_history and len(self.chat_history[model_name]) >= self.MAX_CONVERSATIONS:
            self.error_occurred.emit(f"对话数量已达到上限({self.MAX_CONVERSATIONS})，请先删除一些旧对话！")
            return None
            
        # 初始化模型的对话管理数据结构
        if model_name not in self.chat_history:
            self.chat_history[model_name] = {}
            self.conversation_names[model_name] = {}
            self.next_conversation_id[model_name] = 1
            
        # 创建新对话
        conversation_id = self.next_conversation_id[model_name]
        self.next_conversation_id[model_name] += 1
        
        # 初始对话名称为"新对话"
        self.chat_history[model_name][conversation_id] = []
        self.conversation_names[model_name][conversation_id] = "新对话"
        
        # 设置为当前对话
        self.current_conversation[model_name] = conversation_id
        
        self.save_chat_history()
        self.conversation_updated.emit()
        
        return conversation_id
        
    def rename_conversation(self, model_name, conversation_id, new_name):
        """重命名对话"""
        if not model_name or conversation_id is None:
            return False
            
        if model_name not in self.conversation_names or \
           conversation_id not in self.conversation_names[model_name]:
            return False
            
        self.conversation_names[model_name][conversation_id] = new_name.strip()
        self.save_chat_history()
        self.conversation_updated.emit()
        return True
        
    def delete_conversation(self, model_name, conversation_id):
        """删除对话"""
        if not model_name or conversation_id is None:
            return False
            
        if model_name not in self.chat_history or \
           conversation_id not in self.chat_history[model_name]:
            return False
            
        # 删除对话
        del self.chat_history[model_name][conversation_id]
        del self.conversation_names[model_name][conversation_id]
        
        # 如果删除的是当前对话，需要更新当前对话ID
        if self.current_conversation.get(model_name) == conversation_id:
            if self.chat_history[model_name]:
                self.current_conversation[model_name] = max(self.chat_history[model_name].keys())
            else:
                self.current_conversation[model_name] = None
        
        self.save_chat_history()
        self.conversation_updated.emit()
        return True
        
    def add_message(self, model_name, message, role="user"):
        """添加消息到当前对话"""
        if not model_name:
            return False
            
        conversation_id = self.current_conversation.get(model_name)
        if conversation_id is None:
            return False
            
        if model_name not in self.chat_history or \
           conversation_id not in self.chat_history[model_name]:
            return False
            
        # 添加消息
        self.chat_history[model_name][conversation_id].append({
            "role": role,
            "content": message
        })
        
        # 如果是用户的第一条消息，且当前对话名称是"新对话"，则更新对话名称
        if (role == "user" and 
            len(self.chat_history[model_name][conversation_id]) == 1 and 
            self.conversation_names[model_name][conversation_id] == "新对话"):
            # 使用消息的前20个字符作为对话名称
            new_name = message[:20].strip()
            if new_name:  # 确保名称不为空
                self.conversation_names[model_name][conversation_id] = new_name
                self.conversation_updated.emit()
        
        self.save_chat_history()
        return True
        
    def get_current_conversation_messages(self, model_name):
        """获取当前对话的所有消息"""
        if not model_name:
            return []
            
        conversation_id = self.current_conversation.get(model_name)
        if conversation_id is None:
            return []
            
        if model_name not in self.chat_history or \
           conversation_id not in self.chat_history[model_name]:
            return []
            
        return self.chat_history[model_name][conversation_id]
        
    def get_conversation_list(self, model_name):
        """获取模型的所有对话"""
        if not model_name or model_name not in self.conversation_names:
            return []
            
        return [(conv_id, name) for conv_id, name in self.conversation_names[model_name].items()]
        
    def export_conversations(self, model_name, file_path=None):
        """导出当前模型的所有对话"""
        if not model_name or model_name not in self.chat_history:
            self.error_occurred.emit("当前模型没有对话记录！")
            return False
            
        if not file_path:
            # 获取当前对话名称用于生成默认文件名
            conversation_id = self.current_conversation.get(model_name)
            if conversation_id is None:
                self.error_occurred.emit("请先选择一个对话！")
                return False
                
            conversation_name = self.conversation_names[model_name][conversation_id]
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            default_filename = f"{model_name}_{conversation_name}_{timestamp}.json"
            default_filename = re.sub(r'[<>:"/\\|?*]', '_', default_filename)
            
            file_path = default_filename
            
        try:
            export_data = {
                "model": model_name,
                "conversations": self.chat_history[model_name],
                "names": self.conversation_names[model_name],
                "next_id": self.next_conversation_id[model_name]
            }
            
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(export_data, f, ensure_ascii=False, indent=4)
                
            return True
        except Exception as e:
            self.error_occurred.emit(f"导出对话时发生错误：{str(e)}")
            return False
            
    def import_conversations(self, file_path, model_exists_checker):
        """导入对话记录"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                import_data = json.load(f)
                
            model_name = import_data["model"]
            
            # 确保模型存在
            if not model_exists_checker(model_name):
                self.error_occurred.emit(f"模型 {model_name} 不存在，请先添加该模型！")
                return False
            
            # 初始化模型的对话管理数据结构
            if model_name not in self.chat_history:
                self.chat_history[model_name] = {}
                self.conversation_names[model_name] = {}
                self.next_conversation_id[model_name] = 1
            
            # 合并对话记录
            self.chat_history[model_name].update(import_data["conversations"])
            self.conversation_names[model_name].update(import_data["names"])
            self.next_conversation_id[model_name] = max(
                self.next_conversation_id[model_name],
                import_data["next_id"]
            )
            
            self.save_chat_history()
            self.conversation_updated.emit()
            return True
            
        except Exception as e:
            self.error_occurred.emit(f"导入对话时发生错误：{str(e)}")
            return False
            
    def save_chat_history(self):
        """保存对话历史到文件"""
        try:
            save_data = {
                "chat_history": self.chat_history,
                "conversation_names": self.conversation_names,
                "next_conversation_id": self.next_conversation_id
            }
            
            with open('chat_history.json', 'w', encoding='utf-8') as f:
                json.dump(save_data, f, ensure_ascii=False, indent=4)
        except Exception as e:
            self.error_occurred.emit(f"保存对话历史时发生错误：{str(e)}")

    def load_chat_history(self):
        """从文件加载对话历史"""
        try:
            if os.path.exists('chat_history.json'):
                with open('chat_history.json', 'r', encoding='utf-8') as f:
                    save_data = json.load(f)
                    
                self.chat_history = save_data.get("chat_history", {})
                self.conversation_names = save_data.get("conversation_names", {})
                self.next_conversation_id = save_data.get("next_conversation_id", {})
        except Exception as e:
            self.chat_history = {}
            self.conversation_names = {}
            self.next_conversation_id = {}

    def set_current_conversation(self, model_name, conversation_id):
        """设置当前对话"""
        if not model_name or conversation_id is None:
            return False
            
        if model_name not in self.chat_history or \
           conversation_id not in self.chat_history[model_name]:
            return False
            
        self.current_conversation[model_name] = conversation_id
        return True 